from robot import Robot
import numpy as np
import math
from SE3 import Pose, matrix2quat

class KinematicsSolver:
    def __init__(self, robot: Robot):
        self.robot = robot
        self.a_list = robot.dh_a_list
        self.d_list = robot.dh_d_list
        
    def forward_kinematics(self, joints_list):
        # simple kinematics model
        tf_matrix = np.diag([1.0, 1.0, 1.0, 1.0])
        robot_tf_tree = self.robot.get_robot_tf_tree(joints_list)
        for i in range(self.robot.dof):
            tf_matrix = np.dot(tf_matrix, robot_tf_tree[i])
        x, y, z = tf_matrix[0:3, 3]
        qx, qy, qz, qw = matrix2quat(tf_matrix)
        pose = Pose(x, y, z, qx, qy, qz, qw)
        return pose
    
    def inverse_kinematics(self, target: Pose):
        '''use analytical solutions to solve 6-dof inverse kinematics problem
        '''
        all_joints = []
        tf_matrix = target.pose_to_matrix()

        nx, ox, ax, px = tf_matrix[0]
        ny, oy, ay, py = tf_matrix[1]
        nz, oz, az, pz = tf_matrix[2]
        a0, a1, a2, a3, a4, a5 = self.a_list
        d1, d2, d3, d4, d5, d6 = self.d_list
        ###################### Analytical solutions #######################
        ##  init
        possible_theta1 = []
        possible_theta2 = []
        possible_theta3 = []
        possible_theta4 = []
        possible_theta5 = []
        possible_theta6 = []
        ### theta1 
        possible_theta1.append(-np.arctan2(-py, px)) # theta1 have 2 possible results
        possible_theta1.append(-np.arctan2(-py, px) + np.pi)
        # print(np.rad2deg(possible_theta1))
        ##########################################################
        ### theta3
        temp = [None, None] 
        A = [None, None] 
        x = a2 * d4
        y = a2 * a3
        z = a2 * a2 + a3 * a3 + d4 * d4
        for i in range(2):   # theta3 have 4 possible results
            temp[i] = a1 - px * np.cos(possible_theta1[i]) - py * np.sin(possible_theta1[i])
            A[i] = z - pz * pz - temp[i] * temp[i]
            # print("possible_theta1 = ", possible_theta1[i])
            # print("temp = ", temp[i])
            # print("A = ", A[i])
            possible_theta3.append(-np.arctan2(-y, x) + np.arctan2(A[i], np.sqrt((4 * x * x + 4 * y * y - A[i] * A[i]))))
            possible_theta3.append(-np.arctan2(-y, x) + np.arctan2(A[i], -np.sqrt((4 * x * x + 4 * y * y - A[i] * A[i]))))
        # print(np.rad2deg(possible_theta3))
        ##########################################################
        ### theta2
        m = [None, None, None, None] 
        n = [None, None, None, None] 
        for i in range(len(possible_theta3)): # theta2 have 4 possible results
            m[i] = a2 + a3 * np.cos(possible_theta3[i]) - d4 * np.sin(possible_theta3[i])
            n[i] = a3 * np.sin(possible_theta3[i]) + d4 * np.cos(possible_theta3[i])
            possible_theta2.append(np.arctan2(-pz * m[i] + n[i] * temp[i//2],-pz * n[i] - m[i] * temp[i//2]))
        # print(np.rad2deg(possible_theta2))
        ##########################################################
        ### theta5
        for i in range(4):
            m[i] = (ax * np.cos(possible_theta1[i//2]) * np.cos(possible_theta2[i]) + ay * np.sin(possible_theta1[i//2]) * np.cos(possible_theta2[i]) + az * (-1) * np.sin(possible_theta2[i]))
            n[i] = -(ax * np.cos(possible_theta1[i//2]) * np.sin(possible_theta2[i]) + ay * np.sin(possible_theta1[i//2]) * np.sin(possible_theta2[i]) - az * (-1) * np.cos(possible_theta2[i]))
        for i in range(4): # theta5 have 8 possible results
            x = n[i] * np.cos(possible_theta3[i]) - m[i] * np.sin(possible_theta3[i])
            y = np.sqrt(pow((ay * np.cos(possible_theta1[i//2]) - ax * np.sin(possible_theta1[i//2])),2) + pow((m[i] * np.cos(possible_theta3[i]) + n[i] * np.sin(possible_theta3[i])),2))
            possible_theta5.append(np.arctan2(y, x))
            possible_theta5.append(np.arctan2(-y, x))
        # print(np.rad2deg(possible_theta5))
        ##########################################################
        ### theta4
        for i in range(8):  # theta4 have 8 possible results
            if (np.fabs(np.sin(possible_theta5[i])) >= 1e-4):
                y = ((ay*np.cos(possible_theta1[i//4])-ax*np.sin(possible_theta1[i//4]))) / np.sin(possible_theta5[i])
                x = (-m[i//2]*np.cos(possible_theta3[i//2])-n[i//2]*np.sin(possible_theta3[i//2]))/(np.sin(possible_theta5[i]))
                possible_theta4.append(np.arctan2(y, x))
            else:
                possible_theta4.append(0)
        # print(np.rad2deg(possible_theta4))
        ##########################################################
        ##  theta6 
        for i in range(8):    # theta6 have 8 possible results
            u = nx * np.sin(possible_theta1[i//4]) - ny * np.cos(possible_theta1[i//4])
            v = ox * np.sin(possible_theta1[i//4]) - oy * np.cos(possible_theta1[i//4])
            possible_theta6.append(np.arctan2(np.cos(possible_theta4[i]) * u - np.cos(possible_theta5[i]) * np.sin(possible_theta4[i]) * v, 
                                            np.cos(possible_theta4[i]) * v + np.cos(possible_theta5[i]) * np.sin(possible_theta4[i]) * u))
        # print(np.rad2deg(possible_theta6))
        ##########################################################
        for i in range(8):
            all_joints.append([possible_theta1[i//4],
                               possible_theta2[i//2],
                               possible_theta3[i//2],
                               possible_theta4[i],
                               possible_theta5[i],
                               possible_theta6[i]])
        return all_joints
    
    def compute_best_ik(self, all_joints, start_joint = None):
        def find_max_joint_diff(end_joints, init_joint):
            max = -1000
            for i in range(6):
                if(np.fabs(end_joints[i] - init_joint[i]) > max):
                    max = np.fabs(end_joints[i] - init_joint[i])
            return max
        
        flag = [0, 0, 0, 0, 0, 0, 0, 0]
        ik_result = {"result":[], "success":False}
        if start_joint == None:
            start_joint = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
        ## check joint whether whthin joint limit
        for i in range(8):
            if (np.isnan(all_joints[i][0]) or np.isnan(all_joints[i][1]) or np.isnan(all_joints[i][2]) or
                np.isnan(all_joints[i][3]) or np.isnan(all_joints[i][4]) or np.isnan(all_joints[i][5])):
                continue;
            if(((all_joints[i][0] > self.robot.joint1_limit_min) and (all_joints[i][0] < self.robot.joint1_limit_max)) and 
               ((all_joints[i][1] > self.robot.joint2_limit_min) and (all_joints[i][1] < self.robot.joint2_limit_max)) and 
               ((all_joints[i][2] > self.robot.joint3_limit_min) and (all_joints[i][2] < self.robot.joint3_limit_max)) and 
               ((all_joints[i][3] > self.robot.joint4_limit_min) and (all_joints[i][3] < self.robot.joint4_limit_max)) and 
               ((all_joints[i][4] > self.robot.joint5_limit_min) and (all_joints[i][4] < self.robot.joint5_limit_max)) and 
               ((all_joints[i][5] > self.robot.joint6_limit_min) and (all_joints[i][5] < self.robot.joint6_limit_max))):
                flag[i] = 1;
                ik_result["success"] = True
        best_ik_index = -1
        min_max_joint_angle = 1000;
        if(ik_result["success"]):
            # We calculate the maximum joint rotation value for each IK result (temp_max), 
            # where the minimum value among is the optimal solution index.
            for i in range(8):
                if not flag[i]:
                    continue
                temp_max = find_max_joint_diff(all_joints[i], start_joint)
                if(temp_max < min_max_joint_angle):
                    min_max_joint_angle = temp_max
                    best_ik_index = i
            ik_result["result"] = all_joints[best_ik_index]
        return ik_result
    
    def get_best_ik(self, target_pose:Pose, start_joint = None):
        all_ik = self.inverse_kinematics(target_pose)
        best_ik_result = self.compute_best_ik(all_ik, start_joint)
        return best_ik_result
        